/*
 * =====================================================================================
 *
 *       Filename:  pipeline.c
 *
 *    Description:  Final optimized pipeline implementation with asynchronous inference
 *                  and display, using a ping-pong buffer mechanism for smooth video.
 *
 * =====================================================================================
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>
#include <time.h>

#include "pipeline.h"
#include "util.h"

// --- 全局变量用于推理和显示的异步协作 (乒乓缓冲机制) ---
static pthread_mutex_t g_result_mutex = PTHREAD_MUTEX_INITIALIZER;
static detect_result_group_t g_latest_results; // 用于存储最新的检测结果
static volatile bool g_has_new_result = false;  // 标记是否有新的结果可供绘制

// --- FPS Counter ---
static void print_fps()
{
    static struct timespec start_time = {0};
    static int frame_count = 0;
    const int interval = 2; // 每2秒打印一次

    if (start_time.tv_sec == 0)
    {
        clock_gettime(CLOCK_MONOTONIC, &start_time);
    }

    frame_count++;
    struct timespec current_time;
    clock_gettime(CLOCK_MONOTONIC, &current_time);

    double elapsed = (current_time.tv_sec - start_time.tv_sec) +
                     (current_time.tv_nsec - start_time.tv_nsec) / 1e9;

    if (elapsed >= interval)
    {
        printf("FPS: %.2f (Display)\n", frame_count / elapsed);
        frame_count = 0;
        start_time = current_time;
    }
}

// --- Frame Processing Functions ---

/**
 * @brief 在新帧上执行RGA预处理和RKNN推理，并将结果更新到全局
 */
static void inference_frame(pipeline_context_t *ctx, int dma_fd) {
    // 1. RGA Pre-processing
    if (rga_preprocess_frame(ctx->rga, dma_fd, ctx->v4l2->width, ctx->v4l2->height, ctx->v4l2->in_fourcc) != 0) {
        fprintf(stderr, "ERROR: Frame preprocessing failed.\n");
        return;
    }

    // 2. RKNN Inference
    detect_result_group_t current_results;
    memset(&current_results, 0, sizeof(current_results));
    if (rknn_engine_run(ctx->rknn, ctx->rga->npu_input_mem, &current_results, ctx->v4l2->width, ctx->v4l2->height) != 0) {
        fprintf(stderr, "ERROR: RKNN inference failed.\n");
    }

    // 3. 安全地更新全局的最新结果
    pthread_mutex_lock(&g_result_mutex);
    g_latest_results = current_results;
    g_has_new_result = true;
    pthread_mutex_unlock(&g_result_mutex);
}

/**
 * @brief 在准备显示的帧上绘制最新的检测结果
 */
static void draw_on_frame(pipeline_context_t *ctx, int dma_fd) {
    pthread_mutex_lock(&g_result_mutex);
    // 只要有结果（无论是刚生成的还是上一帧的），就绘制它
    if (g_has_new_result && g_latest_results.count > 0) {
        rga_draw_results(dma_fd, ctx->v4l2->width, ctx->v4l2->height, &g_latest_results);
    }
    pthread_mutex_unlock(&g_result_mutex);
}


// --- Thread Functions (Final Asynchronous Architecture) ---

/**
 * @brief 推理线程: 循环地从V4L2获取帧，执行完整的处理（推理+画框），然后将其交给显示线程
 */
// 在 pipeline.c 中，用这个更简洁的版本替换 inference_thread_loop
static void *inference_thread_loop(void *arg)
{
    pipeline_context_t *ctx = (pipeline_context_t *)arg;
    printf("Inference thread started.\n");

    while (!*ctx->quit_flag)
    {
        struct v4l2_buffer v4l2_buf;
        struct v4l2_plane planes[VIDEO_MAX_PLANES];
        memset(&v4l2_buf, 0, sizeof(v4l2_buf));
        v4l2_buf.m.planes = planes;
        v4l2_buf.length = ctx->v4l2->num_planes;

        // 从V4L2取出一帧
        if (v4l2_dequeue_buffer(ctx->v4l2, &v4l2_buf) < 0) {
            if (errno != EAGAIN && errno != EINTR) {
                fprintf(stderr, "Inference thread: v4l2_dequeue_buffer failed: %s\n", strerror(errno));
                *ctx->quit_flag = true;
            }
            break;
        }

        // --- 在新缓冲区上执行耗时操作 ---
        inference_frame(ctx, v4l2_buf.m.planes[0].m.fd);
        draw_on_frame(ctx, v4l2_buf.m.planes[0].m.fd);
        
        // --- 通知显示线程有新的、已经绘制好的帧可以使用 ---
        pthread_mutex_lock(&ctx->display_mutex);
        ctx->latest_idx = v4l2_buf.index;
        pthread_cond_signal(&ctx->display_cond);
        pthread_mutex_unlock(&ctx->display_mutex);
    }
    
    // 唤醒显示线程，使其能够退出
    pthread_cond_signal(&ctx->display_cond);
    printf("Inference thread finished.\n");
    return NULL;
}


// 在 pipeline.c 中，用这个最终修复版替换 display_thread_loop
static void *display_thread_loop(void *arg)
{
    pipeline_context_t *ctx = (pipeline_context_t *)arg;
    printf("Display thread started (Final Corrected Mode).\n");

    // --- 等待并设置第一帧 ---
    pthread_mutex_lock(&ctx->display_mutex);
    while (ctx->latest_idx == -1 && !*ctx->quit_flag) {
        pthread_cond_wait(&ctx->display_cond, &ctx->display_mutex);
    }
    int idx_to_display = ctx->latest_idx;
    ctx->display_idx = idx_to_display;
    pthread_mutex_unlock(&ctx->display_mutex);

    if (*ctx->quit_flag) return NULL;
    
    // 同步设置第一帧
    drm_set_plane(ctx->drm, idx_to_display);
    
    // --- 主显示循环 ---
    while (!*ctx->quit_flag)
    {
        int prev_display_idx;

        pthread_mutex_lock(&ctx->display_mutex);
        while (ctx->latest_idx == ctx->display_idx && !*ctx->quit_flag) {
            pthread_cond_wait(&ctx->display_cond, &ctx->display_mutex);
        }
        if (*ctx->quit_flag) {
            pthread_mutex_unlock(&ctx->display_mutex);
            break;
        }
        
        // 准备显示最新的一帧
        idx_to_display = ctx->latest_idx;
        pthread_mutex_unlock(&ctx->display_mutex);
        
        // 请求翻页到这个新帧
        if (drm_page_flip(ctx->drm, idx_to_display) < 0) {
            fprintf(stderr, "Display thread: drm_page_flip failed.\n");
            *ctx->quit_flag = true;
            break;
        }
        
        prev_display_idx = ctx->display_idx; // 记录下即将被替换掉的帧
        
        // 更新当前正在显示的缓冲区索引
        pthread_mutex_lock(&ctx->display_mutex);
        ctx->display_idx = idx_to_display;
        pthread_mutex_unlock(&ctx->display_mutex);
        
        // 等待本次翻页完成
        drm_wait_for_flip(ctx->drm);

        // **关键修复**: 翻页完成后，归还上一帧 (prev_display_idx)
        if (prev_display_idx != -1) {
             struct v4l2_buffer temp_buf_to_queue;
             struct v4l2_plane plane;
             memset(&temp_buf_to_queue, 0, sizeof(temp_buf_to_queue));
             temp_buf_to_queue.type = V4L2_BUF_TYPE_VIDEO_CAPTURE_MPLANE;
             temp_buf_to_queue.memory = V4L2_MEMORY_DMABUF;
             temp_buf_to_queue.index = prev_display_idx;
             temp_buf_to_queue.m.planes = &plane;
             temp_buf_to_queue.length = 1;
             temp_buf_to_queue.m.planes[0].m.fd = ctx->drm->buffers[prev_display_idx].dma_fd;
             if (v4l2_queue_buffer(ctx->v4l2, &temp_buf_to_queue) < 0) {
                 fprintf(stderr, "Display thread: v4l2_queue_buffer failed.\n");
                 *ctx->quit_flag = true;
                 break;
             }
        }

        print_fps();
    }
    
    // 退出前，尝试归还最后一帧
    if(ctx->display_idx != -1) {
        // ... (省略重复的归还代码，与上面类似)
    }

    printf("Display thread finished.\n");
    return NULL;
}


// --- Public API ---

int pipeline_init(pipeline_context_t *ctx, v4l2_context_t *v4l2_ctx, drm_device_t *drm_ctx,
                  rknn_engine_context_t *rknn_ctx, rga_context_t *rga_ctx, volatile bool *quit_flag)
{
    ctx->v4l2 = v4l2_ctx;
    ctx->drm = drm_ctx;
    ctx->rknn = rknn_ctx;
    ctx->rga = rga_ctx;
    ctx->quit_flag = quit_flag;

    // 初始化新的同步原语
    pthread_mutex_init(&ctx->display_mutex, NULL);
    pthread_cond_init(&ctx->display_cond, NULL);
    ctx->latest_idx = -1;
    ctx->display_idx = -1;

    return 0;
}

void pipeline_cleanup(pipeline_context_t *ctx)
{
    pthread_mutex_destroy(&ctx->display_mutex);
    pthread_cond_destroy(&ctx->display_cond);
}

int pipeline_start(pipeline_context_t *ctx)
{
    printf("Starting pipeline threads...\n");
    // 改名为 inference_tid
    int ret = pthread_create(&ctx->inference_tid, NULL, inference_thread_loop, ctx);
    RETURN_ON_FAIL(ret != 0, "Failed to create inference thread: %s\n", strerror(ret));

    ret = pthread_create(&ctx->display_tid, NULL, display_thread_loop, ctx);
    if (ret != 0)
    {
        fprintf(stderr, "Failed to create display thread: %s\n", strerror(ret));
        *ctx->quit_flag = true;
        pthread_join(ctx->inference_tid, NULL);
        return -1;
    }
    return 0;
}

void pipeline_stop(pipeline_context_t *ctx)
{
    printf("Stopping pipeline threads...\n");
    
    // 唤醒可能在等待的线程
    pthread_cond_broadcast(&ctx->display_cond);

    if (ctx->inference_tid)
    {
        pthread_join(ctx->inference_tid, NULL);
        ctx->inference_tid = 0;
    }
    if (ctx->display_tid)
    {
        pthread_join(ctx->display_tid, NULL);
        ctx->display_tid = 0;
    }
    printf("Pipeline threads stopped.\n");
}
